from opentelemetry import trace
from fastapi import FastAPI
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.logging import LoggingInstrumentor
import os


class JaegerService:
    """Jaeger分布式追踪工具类
    
    提供FastAPI应用集成Jaeger分布式追踪的工具方法
    """
    
    @staticmethod
    def setup_jaeger(app: FastAPI, service_name: str, port: int):
        """为FastAPI应用设置Jaeger分布式追踪
        
        Args:
            app: 外部传入的FastAPI应用实例
            service_name: 服务名称，将在Jaeger UI中显示
            port: 服务监听端口号
        """
        # 配置服务资源标识
        resource = Resource.create(
            {SERVICE_NAME: service_name}
        )

        # 初始化追踪器
        tracer_provider = TracerProvider(resource=resource)
        trace.set_tracer_provider(tracer_provider)

        # 配置Jaeger导出器，从环境变量获取地址，默认使用jaeger服务名（与docker-compose.yml中的容器名保持一致）
        jaeger_host = os.environ.get('JAEGER_HOST', 'jaeger')
        jaeger_port = int(os.environ.get('JAEGER_PORT', '6831'))
        
        print(f"Setting up Jaeger with host: {jaeger_host}, port: {jaeger_port}")
        
        jaeger_exporter = JaegerExporter(
            agent_host_name=jaeger_host,
            agent_port=jaeger_port,
        )

        # 设置批量数据处理器
        span_processor = BatchSpanProcessor(jaeger_exporter)
        tracer_provider.add_span_processor(span_processor)

        # 为FastAPI应用添加追踪功能
        FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)

        # 为日志添加追踪上下文
        LoggingInstrumentor().instrument(set_logging_format=True)

        # 检查是否已存在相同路径的端点
        existing_paths = {route.path for route in app.routes}
        
        # 添加健康检查路由
        if "/health" not in existing_paths:
            @app.get("/health", tags=["Health"])
            async def health_check():
                return {"status": "UP", "service": service_name}

        # 添加测试路由（用于验证追踪功能）
        if "/test" not in existing_paths:
            @app.get("/test", tags=["Test"])
            async def test_endpoint():
                # 获取当前span并添加一些属性
                current_span = trace.get_current_span()
                current_span.set_attribute("test.attribute", "test-value")
                current_span.add_event(f"{service_name} Test endpoint called")
                current_span.set_attribute("server.port", port)
                return {"message": "Test endpoint reached successfully", "trace_id": current_span.get_span_context().trace_id}